"""
Reve API Client Module

This module provides a comprehensive HTTP client for interacting with the Reve API,
supporting Create, Edit, and Remix operations with proper authentication, error handling,
and retry logic.
"""

import os
import json
import time
import base64
import logging
from typing import Dict, Any, Optional, Union, List
from dataclasses import dataclass
from enum import Enum

try:
    import requests
    from requests.adapters import HTTPAdapter
    from urllib3.util.retry import Retry
except ImportError:
    raise ImportError("requests library is required. Install with: pip install requests")


# Configure logging
logger = logging.getLogger(__name__)





class ReveOperation(Enum):
    """Supported Reve API operations"""
    CREATE = "create"
    EDIT = "edit"
    REMIX = "remix"


class ReveError(Exception):
    """Base exception for Reve API errors"""
    def __init__(self, message: str, error_code: str = None, request_id: str = None):
        super().__init__(message)
        self.error_code = error_code
        self.request_id = request_id


class ReveAuthenticationError(ReveError):
    """Authentication related errors"""
    pass


class ReveContentPolicyError(ReveError):
    """Content policy violation errors"""
    pass


class ReveRateLimitError(ReveError):
    """Rate limiting errors"""
    def __init__(self, message: str, retry_after: int = None, **kwargs):
        super().__init__(message, **kwargs)
        self.retry_after = retry_after


@dataclass
class ReveResponse:
    """Standardized response from Reve API"""
    success: bool
    image: Optional[str] = None  # Base64 encoded image
    metadata: Optional[Dict[str, Any]] = None
    request_id: Optional[str] = None
    credits_used: Optional[int] = None
    credits_remaining: Optional[int] = None
    error_code: Optional[str] = None
    error_message: Optional[str] = None


class ReveAPIClient:
    """
    Core HTTP client for Reve API operations.
    
    Handles authentication, request formatting, response parsing,
    error handling, and retry logic for all Reve API endpoints.
    """
    
    BASE_URL = "https://api.reve.com"
    DEFAULT_TIMEOUT = 60
    MAX_RETRIES = 3
    BACKOFF_FACTOR = 1.0
    
    def __init__(self, api_key: str, timeout: int = DEFAULT_TIMEOUT):
        """
        Initialize the Reve API client.

        Args:
            api_key: Reve API key (required)
            timeout: Request timeout in seconds
        """
        if not api_key or api_key.strip() == "":
            raise ReveAuthenticationError("API key is required.")

        self.api_key = api_key.strip()
        self.timeout = timeout
        self.session = self._create_session()
        
    def _create_session(self) -> requests.Session:
        """Create a requests session with retry strategy and authentication."""
        session = requests.Session()
        
        # Configure retry strategy
        retry_strategy = Retry(
            total=self.MAX_RETRIES,
            backoff_factor=self.BACKOFF_FACTOR,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["POST"]
        )
        
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("http://", adapter)
        session.mount("https://", adapter)
        
        # Set default headers
        session.headers.update({
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "User-Agent": "ComfyUI-Reve-API/1.0"
        })
        
        return session
    
    def _make_request(self, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Make HTTP request to Reve API endpoint.

        Args:
            endpoint: API endpoint path
            payload: Request payload

        Returns:
            Parsed JSON response

        Raises:
            ReveError: For various API errors
        """
        url = f"{self.BASE_URL}{endpoint}"

        try:
            logger.debug(f"Making request to {url}")
            response = self.session.post(url, json=payload, timeout=self.timeout)

            # Handle rate limiting with custom retry
            if response.status_code == 429:
                retry_after = int(response.headers.get("Retry-After", 60))
                raise ReveRateLimitError(
                    f"Rate limit exceeded. Retry after {retry_after} seconds.",
                    retry_after=retry_after,
                    request_id=response.headers.get("X-Request-ID")
                )

            # Check if we got HTML instead of JSON (API not available)
            content_type = response.headers.get("content-type", "").lower()
            if "text/html" in content_type:
                raise ReveError(
                    "Reve API endpoints are not available. The API may be in development or requires special access. "
                    "Please check https://api.reve.com for the latest API status or contact Reve support.",
                    "API_NOT_AVAILABLE",
                    None
                )

            # Parse response
            try:
                response_data = response.json()
            except json.JSONDecodeError:
                # Special handling for 404 errors that return HTML
                if response.status_code == 404:
                    raise ReveError(
                        "Reve API endpoint not found. The API endpoints may have changed or the API is not yet publicly available. "
                        "Please check https://api.reve.com for the latest API documentation.",
                        "ENDPOINT_NOT_FOUND",
                        None
                    )
                response_data = {"error": {"message": "Invalid JSON response"}}

            # Handle HTTP errors
            if not response.ok:
                # Try to get detailed error information from response
                error_info = response_data.get("error", {})

                # If no error object, check for direct error fields
                if not error_info:
                    error_info = {
                        "code": response_data.get("code"),
                        "message": response_data.get("message") or response_data.get("error")
                    }

                error_code = error_info.get("code", f"HTTP_{response.status_code}")
                error_message = error_info.get("message", f"HTTP {response.status_code} error")
                request_id = response_data.get("request_id") or response.headers.get("X-Request-ID")

                # Add raw response data for debugging if error message is generic
                if error_message == f"HTTP {response.status_code} error":
                    logger.debug(f"Raw response data: {response_data}")
                    # Try to extract any error information from the response
                    if isinstance(response_data, dict):
                        # Look for common error field names
                        for field in ["detail", "details", "msg", "description"]:
                            if field in response_data:
                                error_message = f"HTTP {response.status_code}: {response_data[field]}"
                                break
                        else:
                            # If we have response data but no clear error message, include it
                            if response_data:
                                error_message = f"HTTP {response.status_code}: {json.dumps(response_data)}"

                if response.status_code == 401:
                    raise ReveAuthenticationError(error_message, error_code, request_id)
                elif response.status_code == 404:
                    raise ReveError(
                        "Reve API endpoint not found. The API endpoints may have changed or the API is not yet publicly available. "
                        "Please check https://api.reve.com for the latest API documentation.",
                        "ENDPOINT_NOT_FOUND",
                        request_id
                    )
                elif error_code == "CONTENT_POLICY_VIOLATION":
                    raise ReveContentPolicyError(error_message, error_code, request_id)
                else:
                    raise ReveError(error_message, error_code, request_id)

            return response_data
            
        except requests.exceptions.Timeout:
            raise ReveError(f"Request timeout after {self.timeout} seconds")
        except requests.exceptions.ConnectionError:
            raise ReveError("Connection error. Please check your internet connection.")
        except requests.exceptions.RequestException as e:
            raise ReveError(f"Request failed: {str(e)}")
    
    def create(self, prompt: str, aspect_ratio: str = "3:2", version: str = "latest") -> ReveResponse:
        """
        Create image from text prompt.

        Args:
            prompt: Text description of desired image
            aspect_ratio: Image aspect ratio (16:9, 9:16, 3:2, 2:3, 4:3, 3:4, 1:1)
            version: Model version (latest, reve-create@20250915)

        Returns:
            ReveResponse with generated image and metadata (always returns PNG format)
        """
        payload = {
            "prompt": prompt,
            "aspect_ratio": aspect_ratio,
            "version": version
        }
        
        try:
            response_data = self._make_request("/v1/image/create", payload)

            # Create metadata from response fields
            metadata = {
                "version": response_data.get("version"),
                "content_violation": response_data.get("content_violation"),
                "aspect_ratio": aspect_ratio,
                "prompt": prompt
            }

            return ReveResponse(
                success=True,
                image=response_data.get("image"),
                metadata=metadata,
                request_id=response_data.get("request_id"),
                credits_used=response_data.get("credits_used"),
                credits_remaining=response_data.get("credits_remaining")
            )
        except ReveError as e:
            return ReveResponse(
                success=False,
                error_code=e.error_code,
                error_message=str(e),
                request_id=e.request_id
            )
    
    def edit(self, prompt: str, reference_image: str, version: str = "latest") -> ReveResponse:
        """
        Edit existing image with text instructions.

        Args:
            prompt: Text instruction for editing (sent as edit_instruction to API)
            reference_image: Base64 encoded source image
            version: Model version (latest, reve-edit@20250915)


        Returns:
            ReveResponse with edited image and metadata
        """
        payload = {
            "edit_instruction": prompt,  # API expects 'edit_instruction', not 'prompt'
            "reference_image": reference_image,
            "version": version
        }
        
        try:
            response_data = self._make_request("/v1/image/edit", payload)

            # Create metadata from response fields
            metadata = {
                "version": response_data.get("version"),
                "content_violation": response_data.get("content_violation"),
                "prompt": prompt,
                "reference_image": reference_image
            }

            return ReveResponse(
                success=True,
                image=response_data.get("image"),
                metadata=metadata,
                request_id=response_data.get("request_id"),
                credits_used=response_data.get("credits_used"),
                credits_remaining=response_data.get("credits_remaining")
            )
        except ReveError as e:
            return ReveResponse(
                success=False,
                error_code=e.error_code,
                error_message=str(e),
                request_id=e.request_id
            )
    
    def remix(self, prompt: str, reference_images: List[str], aspect_ratio: str = "3:2",
              version: str = "latest") -> ReveResponse:
        """
        Remix multiple images with text guidance.
        
        Args:
            prompt: Text description with XML tags for image references
            reference_images: List of 1-4 base64 encoded images
            aspect_ratio: Output image aspect ratio
            version: Model version (latest, reve-remix@20250915)

            
        Returns:
            ReveResponse with remixed image and metadata
        """
        if not reference_images or len(reference_images) > 4:
            return ReveResponse(
                success=False,
                error_code="INVALID_PARAMETERS",
                error_message="reference_images must contain 1-4 images"
            )
        
        payload = {
            "prompt": prompt,
            "reference_images": reference_images,
            "aspect_ratio": aspect_ratio,
            "version": version
        }
        
        try:
            response_data = self._make_request("/v1/image/remix", payload)

            # Create metadata from response fields
            metadata = {
                "version": response_data.get("version"),
                "content_violation": response_data.get("content_violation"),
                "aspect_ratio": aspect_ratio,
                "prompt": prompt,
                "reference_images": reference_images
            }

            return ReveResponse(
                success=True,
                image=response_data.get("image"),
                metadata=metadata,
                request_id=response_data.get("request_id"),
                credits_used=response_data.get("credits_used"),
                credits_remaining=response_data.get("credits_remaining")
            )
        except ReveError as e:
            return ReveResponse(
                success=False,
                error_code=e.error_code,
                error_message=str(e),
                request_id=e.request_id
            )
    
    def get_credits(self) -> Dict[str, Any]:
        """
        Get current credit balance (if supported by API).
        
        Returns:
            Dictionary with credit information
        """
        # This endpoint may not exist - implement if available
        try:
            response_data = self._make_request("/api/credits", {})
            return response_data
        except ReveError:
            # Fallback - credits are returned with each operation
            return {"error": "Credits endpoint not available"}
    
    def close(self):
        """Close the HTTP session."""
        if self.session:
            self.session.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


# Utility function for easy client creation
def create_reve_client(api_key: str, timeout: int = 60) -> ReveAPIClient:
    """
    Create a Reve API client instance.

    Args:
        api_key: API key (required)
        timeout: Request timeout in seconds

    Returns:
        Configured ReveAPIClient instance
    """
    return ReveAPIClient(api_key=api_key, timeout=timeout)
